from llmtcg.chat.models.openai import OpenaiCompatibleChatModel

from typing import Any, Dict, Tuple
from os import environ
from dataclasses import dataclass
from enum import Enum


@dataclass
class MistralChatModel(OpenaiCompatibleChatModel):

    """ Models that are deployed at Infini-AI. You need to set the environment
    variable INFINIAI_API_KEY or manually provide the api key when initialize
    the model instance. """

    class Model(Enum):
        OPEN_8_7B = "open-mixtral-8x7b"
        SMALL = "mistral-small-latest"

    api_key: str = environ.get("MISTRAL_API_KEY", "")
    model: Model = Model.OPEN_8_7B

    @property
    def api_address(self) -> str:
        return "https://api.mistral.ai/v1"

    @property
    def headers(self) -> Dict[str, str]:
        return {
            "Content-Type": "application/json",
            "Accept": "application/json",
            "Authorization": f"Bearer {self.api_key}"
        }

    @property
    def model_name(self) -> str:
        return self.model.value

    def detect_error(self, response: Any) -> Tuple[bool, str]:
        if 'code' in response:
            if response['code'] == 10003:
                return True, f"authentication error, please check API_KEY (current value is {self.api_key})"

        return False, ""

    def __hash__(self) -> int:
        return id(self)
